import os
import wave
import pyaudio
import numpy as np
from sphfile import SPHFile


def WAV2wav():
    """ 将 TIMIT 的 WAV 不可读文件转为 wav 文件
    """
    for root, ds, fs in os.walk("./data/TIMIT"):
        for f in fs:
            if f.endswith(".WAV"):
                fullname = os.path.join(root, f)
                dirname = os.path.dirname(fullname).replace("TIMIT", "TIMIT-T")
                if not os.path.exists(dirname):
                    os.makedirs(dirname)
                sph = SPHFile(fullname)
                sph.write_wav(fullname.replace("TIMIT", "TIMIT-T").replace(".WAV", "-t.wav"))


def record(data, filename, p, ch=1, sr=16000):
    """ 将处理过后的 wav 数据写回磁盘中
    """
    wf = wave.open(filename, "wb")
    wf.setnchannels(ch)
    wf.setsampwidth(p.get_sample_size(pyaudio.paInt16))
    wf.setframerate(sr)
    wf.writeframes(data)
    wf.close()


def audio_simultaneous(fl1, fl2, audio_len=None):
    """ 将两个音频文件处理等长后返回

    fl1,fl2: 两个音频文件
    audio_len: 指定音频文件的长度
    """

    # 读取文件
    f1 = wave.open(fl1, 'rb')
    f2 = wave.open(fl2, 'rb')

    # 检查一下音频文件的采样率
    if not f1.getparams()[2] == 16000:
        raise Exception(f"{fl1} sample rate is not 16000")
    if not f2.getparams()[2] == 16000:
        raise Exception(f"{fl2} sample rate is not 16000")
    
    # 读取 wav 数据
    nf1, nf2 = f1.getparams()[3], f2.getparams()[3]
    f1_wav_data = np.frombuffer(f1.readframes(nf1), dtype=np.int16)
    f2_wav_data = np.frombuffer(f2.readframes(nf2), dtype=np.int16)

    # 统一长度
    if audio_len == None:
        length = abs(nf1 - nf2)
        temp_array = np.zeros(length, dtype=np.int16)
        if nf1 > nf2:
            f2_wav_data = np.concatenate((f2_wav_data, temp_array))
        elif nf1 < nf2:
            f1_wav_data = np.concatenate((f1_wav_data, temp_array))
    else:
        f1_wav_data = np.concatenate((f1_wav_data, np.zeros(audio_len - nf1, dtype=np.int16)))
        f2_wav_data = np.concatenate((f2_wav_data, np.zeros(audio_len - nf2, dtype=np.int16)))
    
    # 关闭文件
    f1.close()
    f2.close()

    return f1_wav_data, f2_wav_data


def mix_audio(new_root, audio_len=None):
    """ 混合音频，将混合后的音频存储到

    new_root: 混合后的音频的存储路径
    audio_len: 指定了统一的音频长度，默认不指定
    """

    idx = 1  # 新保存音频文件的 id
    p = pyaudio.PyAudio()  # 申明一个录音机

    # 混合音频
    for rt in ["./data/TIMIT-T/TRAIN/DR{}", "./data/TIMIT-T/TEST/DR{}"]:
        for i in range(1, 9):
            root = rt.format(i)
            names = os.listdir(root)
            for nid in range(1, len(names)):
                pt1 = os.path.join(root, names[nid - 1])
                pt2 = os.path.join(root, names[nid])
                for fl1, fl2 in zip(os.listdir(pt1), os.listdir(pt2)):
                    fl1 = os.path.join(pt1, fl1)
                    fl2 = os.path.join(pt2, fl2)
                    
                    # 建立新音频保存路径
                    save_root = os.path.join(root.replace("TIMIT-T", new_root), "{:04}".format(idx))
                    if not os.path.exists(save_root):
                        os.makedirs(save_root)
                    
                    # 音频处理和保存
                    f1_wav_data, f2_wav_data = audio_simultaneous(fl1, fl2, audio_len=audio_len)
                    record(f1_wav_data.tobytes(), os.path.join(save_root, "{:04}-s1.wav".format(idx)), p)
                    record(f2_wav_data.tobytes(), os.path.join(save_root, "{:04}-s2.wav".format(idx)), p)
                    new_data = f1_wav_data + f2_wav_data  # 混合音频
                    record(new_data.tobytes(), os.path.join(save_root, "{:04}-mix.wav".format(idx)), p)

                    idx += 1


def find_audio_len():
    """ 找到所有音频文件中，数据最长有多长
    最后输出结果: (14644, 124621)
    """
    max_len = 0
    min_len = float('inf')
    for root, ds, fs in os.walk("./data/TIMIT-T"):
        for f in fs:
            if f.endswith(".wav"):
                ad = wave.open(os.path.join(root, f), 'rb')
                l = ad.getparams()[3]
                if l > max_len:
                    max_len = l
                if l < min_len:
                    min_len = l
                ad.close()
    return min_len, max_len


if __name__ == "__main__":
    mix_audio("TIMIT-mix-8", audio_len=128000)
